import os
from create_sub_file.write_pbs.write_pbs_common import *
from create_sub_file.common_file import *

def create_para(tmp_amg_dict):
    run_para = ""
    for i in tmp_amg_dict:
        run_para = run_para + "-{0} {1} ".format(i,tmp_amg_dict[i])
    run_para = run_para.strip(" ")
    return run_para

def write_pbs_mpirun(fo,i,matrix,rhs,mpi_option,hpctoolkit_option,para_dict,other_para,dir):
    amg_para = create_para(para_dict)
    programname = os.path.basename(i)
    fo.write("cd $WORKDIR/{0}\n".format(programname))
    fo.write(f"echo \"====run program {programname}====\"\n\n")
    if len(hpctoolkit_option):
        runstr = "{0} {1} {2} -fromonecsrfile {3} -rhsfromfile {4} {5} {6}\n".format(\
            mpi_option, hpctoolkit_option, i, matrix, rhs, amg_para, other_para)
        if(programname.split("_")[0].lower() == "hypre" ):
            runstr = runstr.replace("-rhsfromfile", "-rhsfromonefile")
    else:
        runstr = "{0} {1} -fromonecsrfile {2} -rhsfromfile {3} {4} {5}\n".format(\
            mpi_option, i, matrix, rhs, amg_para, other_para)
    if(programname.split("_")[0].lower() == "hypre" ):
        runstr = runstr.replace("-rhsfromfile", "-rhsfromonefile")
    fo.write(f"{runstr}\n")
    fo.write(f"echo \"====end program {programname}====\"\n\n")
    return 0

def write_pbs_amg(fo,program,amg_parameter,tmp_amg_dict,runenvironments,dir):
    caseid = tmp_amg_dict["case"]
    problem = amg_parameter["problem"]
    mpi_option = runenvironments["mpi"]["mpirun_option"]
    for i in program:
        programname = os.path.basename(i)
        if programname.split("_")[0].lower() == "jxpamg":
            matrix = problem["jxpamg_matrix_dir"][caseid]
            rhs = problem["rhs_dir"][caseid]
            amg_dict = tmp2run_jxpamg_dict(tmp_amg_dict)
            mpi_option = create_mpi_option(tmp_amg_dict,runenvironments)
            hpctoolkit_option = create_hpctoolkit_option(runenvironments)
            other_para = create_other_para(amg_parameter["run_parameter"])
        elif programname.split("_")[0].lower() == "hypre":
            matrix = problem["hypre_matrix_dir"][caseid]
            rhs = problem["rhs_dir"][caseid]
            amg_dict = tmp2run_hypre_dict(tmp_amg_dict,amg_parameter)
            mpi_option = create_mpi_option(tmp_amg_dict,runenvironments)
            hpctoolkit_option = create_hpctoolkit_option(runenvironments)
            other_para = create_other_para(amg_parameter["HYPRE_corresponding_parameter"])
        write_pbs_mpirun(fo,i,matrix,rhs,mpi_option,hpctoolkit_option,amg_dict,other_para,dir)